package milvus

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"MyBlogv2/common/embedding"
	"MyBlogv2/common/models"

	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"gopkg.in/yaml.v2"
)

// MilvusClient Milvus客户端
type MilvusClient struct {
	client client.Client
	ctx    context.Context
	dbName string
}

var MilvusDB *MilvusClient = InitMilvusClient()

var milvus_config, _ = loadConfig()

func loadConfig() (map[string]string, error) {
	config := map[string]string{}

	// 获取当前工作目录
	workDir, err := os.Getwd()
	if err != nil {
		log.Println("获取工作目录失败:", err)
		return nil, err
	}

	// 尝试多个可能的配置文件路径
	possiblePaths := []string{
		// 如果当前在项目根目录
		filepath.Join(workDir, "common", "milvus", "milvus_config.yaml"),
		// 如果当前在test目录
		filepath.Join(workDir, "..", "common", "milvus", "milvus_config.yaml"),
		// 如果当前在其他子目录
		filepath.Join(workDir, "..", "..", "common", "milvus", "milvus_config.yaml"),
	}

	for _, configPath := range possiblePaths {
		if data, err := ioutil.ReadFile(configPath); err == nil {
			if err := yaml.Unmarshal(data, config); err != nil {
				log.Println("读取配置文件失败", err)
				return nil, err
			}
			log.Println("milvus_config.yaml文件读取成功:", configPath)
			return config, nil
		}
	}

	log.Println("配置文件不存在，尝试的路径:")
	for _, path := range possiblePaths {
		log.Println("  -", path)
	}
	return nil, nil
}

func InitMilvusClient() *MilvusClient {
	client, err := NewMilvusClient()
	if err != nil {
		log.Fatalf("创建Milvus客户端失败: %v", err)
		return nil
	}
	return client
}


// NewMilvusClient 创建Milvus客户端
func NewMilvusClient() (*MilvusClient, error) {
	host := milvus_config["host"]
	portStr := milvus_config["port"]
	dbName := milvus_config["dbName"]

	ctx := context.Background()

	// 将端口字符串转换为整数
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("端口号格式错误: %v", err)
	}

	// 创建连接地址
	addr := fmt.Sprintf("%s:%d", host, port)

	// 创建客户端
	c, err := client.NewGrpcClient(ctx, addr)
	if err != nil {
		return nil, fmt.Errorf("创建Milvus客户端失败: %v", err)
	}

	// 如果指定了数据库，切换到该数据库
	if dbName != "default" {
		err = c.UsingDatabase(ctx, dbName)
		if err != nil {
			return nil, fmt.Errorf("切换到数据库 %s 失败: %v", dbName, err)
		}
	}

	return &MilvusClient{
		client: c,
		ctx:    ctx,
		dbName: dbName,
	}, nil
}

// Close 关闭客户端连接
func (smc *MilvusClient) Close() error {
	return smc.client.Close()
}

// InsertVector 插入向量
func (smc *MilvusClient) InsertVectorToPostCollection(postID int64, vector []float32) error {
	collectionName := milvus_config["collectionName"]
	// 创建post_id列
	postIDColumn := entity.NewColumnInt64("post_id", []int64{postID})

	// 创建向量列
	vectorColumn := entity.NewColumnFloatVector("vector", 2048, [][]float32{vector})

	// 插入数据
	_, err := smc.client.Insert(smc.ctx, collectionName, "", postIDColumn, vectorColumn)
	if err != nil {
		return fmt.Errorf("插入向量数据失败: %v", err)
	}

	log.Printf("成功插入文章ID %d 的向量到集合 %s", postID, collectionName)
	smc.FlushCollection()
	return nil
}

// GetCollectionCount 获取集合中的记录数
func (smc *MilvusClient) GetCollectionCount() (int64, error) {
	smc.FlushCollection()
	collectionName := milvus_config["collectionName"]
	// 获取集合统计信息
	stats, err := smc.client.GetCollectionStatistics(smc.ctx, collectionName)
	if err != nil {
		return 0, fmt.Errorf("获取集合统计信息失败: %v", err)
	}
	fmt.Println(stats)

	// 从统计信息中提取记录数
	for key, value := range stats {
		if key == "row_count" {
			// 将字符串转换为int64
			var count int64
			if _, err := fmt.Sscanf(value, "%d", &count); err == nil {
				return count, nil
			}
		}
	}

	return 0, fmt.Errorf("未找到记录数信息")
}

// GetCountByPostID 通过postId查询数据是否存在，返回匹配的记录数
func (smc *MilvusClient) GetCountByPostID(postID int64) (int64, error) {
	smc.FlushCollection()
	collectionName := milvus_config["collectionName"]
	// 构建查询表达式，查询指定post_id的记录
	expr := fmt.Sprintf("post_id == %d", postID)

	// 执行查询，只返回post_id字段
	resultSet, err := smc.client.Query(smc.ctx, collectionName, nil, expr, []string{"post_id"})
	if err != nil {
		return 0, fmt.Errorf("查询post_id %d 失败: %v", postID, err)
	}

	// 获取post_id列
	postIDColumn := resultSet.GetColumn("post_id")
	if postIDColumn == nil {
		return 0, fmt.Errorf("未找到post_id列")
	}

	// 返回匹配的记录数
	count := int64(postIDColumn.Len())
	log.Printf("post_id %d 在集合 %s 中找到 %d 条记录", postID, collectionName, count)

	return count, nil
}

// DeleteByPostID 通过postId删除数据
func (smc *MilvusClient) DeleteByPostID(postID int64) error {
	collectionName := milvus_config["collectionName"]

	// 构建删除表达式，删除指定post_id的记录
	expr := fmt.Sprintf("post_id == %d", postID)

	// 执行删除操作
	err := smc.client.Delete(smc.ctx, collectionName, "", expr)
	if err != nil {
		return fmt.Errorf("删除post_id %d 失败: %v", postID, err)
	}

	log.Printf("成功删除post_id %d 的数据", postID)
	smc.FlushCollection()
	return nil
}

// UpdateVectorByPostID 根据postId更新vector数据
func (smc *MilvusClient) UpdateVectorByPostID(postID int64, vector []float32) error {
	collectionName := milvus_config["collectionName"]

	// 先删除指定post_id的现有记录
	err := smc.DeleteByPostID(postID)
	if err != nil {
		return fmt.Errorf("删除post_id %d 的现有记录失败: %v", postID, err)
	}

	// 插入新的向量数据
	err = smc.InsertVectorToPostCollection(postID, vector)
	if err != nil {
		return fmt.Errorf("插入post_id %d 的新向量数据失败: %v", postID, err)
	}

	log.Printf("成功更新文章ID %d 的向量到集合 %s", postID, collectionName)
	smc.FlushCollection()
	return nil
}

// SearchResult 搜索结果结构体
type SearchResult struct {
	PostID   int64   `json:"post_id"`
	Score    float32 `json:"score"`
}

// SearchSimilarVectors 相似度查询topk方法
func (smc *MilvusClient) SearchSimilarVectors(queryVector []float32, topK int) ([]SearchResult, error) {
	collectionName := milvus_config["collectionName"]

	// 创建搜索参数，使用IVF_FLAT索引的搜索参数
	sp, err := entity.NewIndexIvfFlatSearchParam(10) // nprobe设置为10
	if err != nil {
		return nil, fmt.Errorf("创建搜索参数失败: %v", err)
	}

	// 执行搜索
	searchResults, err := smc.client.Search(
		smc.ctx,
		collectionName,
		[]string{},          // 分区名称列表，空表示搜索所有分区
		"",                  // 过滤表达式，空表示不过滤
		[]string{"post_id"}, // 输出字段
		[]entity.Vector{entity.FloatVector(queryVector)}, // 查询向量
		"vector",  // 向量字段名称
		entity.L2, // 度量类型：L2表示欧几里得距离
		topK,      // 返回的Top K结果数量
		sp,        // 搜索参数
	)
	if err != nil {
		return nil, fmt.Errorf("相似度搜索失败: %v", err)
	}

	// 处理搜索结果
	var results []SearchResult
	for _, sr := range searchResults {
		// 获取post_id列
		var postIDColumn *entity.ColumnInt64
		for _, field := range sr.Fields {
			if field.Name() == "post_id" {
				if c, ok := field.(*entity.ColumnInt64); ok {
					postIDColumn = c
					break
				}
			}
		}

		if postIDColumn == nil {
			return nil, fmt.Errorf("未找到post_id字段")
		}

		// 提取结果
		for i := 0; i < sr.ResultCount; i++ {
			postID, err := postIDColumn.ValueByIdx(i)
			if err != nil {
				return nil, fmt.Errorf("获取post_id失败: %v", err)
			}

			results = append(results, SearchResult{
				PostID:   postID,
				Score:    sr.Scores[i],
			})
		}
	}

	log.Printf("相似度搜索完成，返回 %d 个结果", len(results))
	return results, nil
}

// SearchSimilarVectorsByText 通过文本进行相似度查询
func (smc *MilvusClient) SearchSimilarVectorsByText(text string, topK int) ([]SearchResult, error) {
	log.Println("查询文本:", text)
	// 生成查询向量
	embeddingResult := embedding.GetEmbeddingByText(text)

	// 解析embedding结果
	var embeddingResp struct {
		Data struct {
			Embedding []float32 `json:"embedding"`
		} `json:"data"`
	}

	err := json.Unmarshal([]byte(embeddingResult), &embeddingResp)
	if err != nil {
		return nil, fmt.Errorf("解析embedding结果失败: %v", err)
	}

	queryVector := embeddingResp.Data.Embedding
	log.Printf("生成查询向量，维度: %d", len(queryVector))

	// 执行相似度搜索
	return smc.SearchSimilarVectors(queryVector, topK)
}

// FlushCollection 刷新集合数据
func (smc *MilvusClient) FlushCollection() error {
	collectionName := milvus_config["collectionName"]
	err := smc.client.Flush(smc.ctx, collectionName, false)
	if err != nil {
		return fmt.Errorf("刷新集合 %s 失败: %v", collectionName, err)
	}
	log.Printf("集合 %s 数据刷新成功", collectionName)
	return nil
}


func (smc *MilvusClient) UpsertArticleToMilvus(article models.Article) error{
	log.Printf("ID=%d, 标题=%s", article.Id, article.ArticleTitle)
	count, err := smc.GetCountByPostID(int64(article.Id))
	if err != nil {
		log.Printf("检查文章ID %d 是否存在时出错: %v", article.Id, err)
		return err
	}

	if count > 0 {
		log.Printf("文章ID %d 已存在于Milvus中，跳过", article.Id)
		return nil
	}

	// 生成向量
	vector, err := generateArticleVector(article)
	if err != nil {
		log.Printf("为文章ID %d 生成向量失败: %v", article.Id, err)
		return err
	}

	// 插入向量到Milvus
	err = smc.InsertVectorToPostCollection(int64(article.Id), vector)
	if err != nil {
		log.Printf("插入文章ID %d 的向量到Milvus失败: %v", article.Id, err)
		return err
	}

	log.Printf("成功导入文章ID %d 的向量", article.Id)
	return nil
}




// generateArticleVector 为文章生成向量
func generateArticleVector(article models.Article) ([]float32, error) {
	// 组合文章的标题和内容作为向量生成的文本
	text := fmt.Sprintf("%s %s", article.ArticleTitle, article.ArticleContent)

	// 清理文本，移除HTML标签和多余空白
	text = cleanText(text)

	// 如果文本太短，使用标题
	if len(text) < 10 {
		text = article.ArticleTitle
	}

	// 如果文本太长，截取前2000个字符
	if len(text) > 2000 {
		text = text[:2000]
	}

	log.Printf("为文章ID %d 生成向量，文本长度: %d", article.Id, len(text))

	// 调用embedding服务生成向量
	embeddingResult := embedding.GetEmbeddingByText(text)

	// 解析embedding结果
	var embeddingResp struct {
		Data struct {
			Embedding []float32 `json:"embedding"`
		} `json:"data"`
	}

	err := json.Unmarshal([]byte(embeddingResult), &embeddingResp)
	if err != nil {
		return nil, fmt.Errorf("解析embedding结果失败: %v", err)
	}

	vector := embeddingResp.Data.Embedding
	log.Printf("文章ID %d 向量生成成功，维度: %d", article.Id, len(vector))

	return vector, nil
}

// cleanText 清理文本，移除HTML标签和多余空白
func cleanText(text string) string {
	// 移除HTML标签（简单实现）
	text = strings.ReplaceAll(text, "<p>", " ")
	text = strings.ReplaceAll(text, "</p>", " ")
	text = strings.ReplaceAll(text, "<br>", " ")
	text = strings.ReplaceAll(text, "<br/>", " ")
	text = strings.ReplaceAll(text, "<br />", " ")
	text = strings.ReplaceAll(text, "<div>", " ")
	text = strings.ReplaceAll(text, "</div>", " ")
	text = strings.ReplaceAll(text, "<span>", " ")
	text = strings.ReplaceAll(text, "</span>", " ")
	text = strings.ReplaceAll(text, "<strong>", " ")
	text = strings.ReplaceAll(text, "</strong>", " ")
	text = strings.ReplaceAll(text, "<em>", " ")
	text = strings.ReplaceAll(text, "</em>", " ")
	text = strings.ReplaceAll(text, "<h1>", " ")
	text = strings.ReplaceAll(text, "</h1>", " ")
	text = strings.ReplaceAll(text, "<h2>", " ")
	text = strings.ReplaceAll(text, "</h2>", " ")
	text = strings.ReplaceAll(text, "<h3>", " ")
	text = strings.ReplaceAll(text, "</h3>", " ")
	text = strings.ReplaceAll(text, "<h4>", " ")
	text = strings.ReplaceAll(text, "</h4>", " ")
	text = strings.ReplaceAll(text, "<h5>", " ")
	text = strings.ReplaceAll(text, "</h5>", " ")
	text = strings.ReplaceAll(text, "<h6>", " ")
	text = strings.ReplaceAll(text, "</h6>", " ")
	text = strings.ReplaceAll(text, "<ul>", " ")
	text = strings.ReplaceAll(text, "</ul>", " ")
	text = strings.ReplaceAll(text, "<ol>", " ")
	text = strings.ReplaceAll(text, "</ol>", " ")
	text = strings.ReplaceAll(text, "<li>", " ")
	text = strings.ReplaceAll(text, "</li>", " ")
	text = strings.ReplaceAll(text, "<a>", " ")
	text = strings.ReplaceAll(text, "</a>", " ")
	text = strings.ReplaceAll(text, "<img>", " ")
	text = strings.ReplaceAll(text, "</img>", " ")
	text = strings.ReplaceAll(text, "<code>", " ")
	text = strings.ReplaceAll(text, "</code>", " ")
	text = strings.ReplaceAll(text, "<pre>", " ")
	text = strings.ReplaceAll(text, "</pre>", " ")

	// 移除多余空白
	text = strings.ReplaceAll(text, "\n", " ")
	text = strings.ReplaceAll(text, "\r", " ")
	text = strings.ReplaceAll(text, "\t", " ")

	// 压缩多个连续空格为单个空格
	for strings.Contains(text, "  ") {
		text = strings.ReplaceAll(text, "  ", " ")
	}

	// 去除首尾空白
	text = strings.TrimSpace(text)

	return text
}




// // CreateVectorIndex 创建向量索引
// func (smc *MilvusClient) CreateVectorIndex(collectionName string) error {
// 	// 创建IVF_FLAT索引
// 	idx, err := entity.NewIndexIvfFlat(entity.L2, 1024)
// 	if err != nil {
// 		return fmt.Errorf("创建索引失败: %v", err)
// 	}

// 	// 为vector字段创建索引
// 	err = smc.client.CreateIndex(smc.ctx, collectionName, "vector", idx, false)
// 	if err != nil {
// 		return fmt.Errorf("创建向量索引失败: %v", err)
// 	}

// 	log.Printf("集合 %s 的向量索引创建成功", collectionName)
// 	return nil
// }

// // DropCollection 删除集合
// func (smc *MilvusClient) DropCollection(collectionName string) error {
// 	err := smc.client.DropCollection(smc.ctx, collectionName)
// 	if err != nil {
// 		return fmt.Errorf("删除集合 %s 失败: %v", collectionName, err)
// 	}

// 	log.Printf("集合 %s 删除成功", collectionName)
// 	return nil
// }
